#include <Eigen/Geometry>

#include "render.hpp"

// get color of `point` in direction `out_dir`
static Eigen::Vector3f shade(
    const Scene & scene,
    int object_idx,
    int face_idx,
    const Eigen::Vector3f & point,
    const Eigen::Vector3f & out_dir,
    const RenderConfig & config)
{
    auto & objects = scene.get_objects();
    auto & object = objects[object_idx];

    Eigen::Vector3f normal, up;
    object.face_tangent(face_idx, normal, up);

    // try hitting lights

    Eigen::Vector3f color = Eigen::Vector3f::Zero();

    for (int i = 0; i < objects.size(); i++)
        if (objects[i].type == 2)
            for (int j = 0; j < objects[i].faces.rows(); j++)
            {
                Eigen::Matrix3f Ps;
                objects[i].face_vertices(j, Ps);

                Eigen::Vector3f E1 = Ps.col(1) - Ps.col(0);
                Eigen::Vector3f E2 = Ps.col(2) - Ps.col(0);

                float b1 = config.rng(), b2 = config.rng();
                if (b1 + b2 > 1.0f)
                {
                    b1 = 1.0f - b1;
                    b2 = 1.0f - b2;
                }

                // a random point on light source
                Eigen::Vector3f sample = Ps.col(0) + E1 * b1 + E2 * b2;
                // reversed incident direction
                Eigen::Vector3f in_dir = (sample - point).normalized();

                float dist;
                int object_idx, face_idx;
                Eigen::Vector3f sky_color;
                scene.intersect(point, in_dir, dist, object_idx, face_idx, sky_color);
                if (object_idx != i || face_idx != j) // sample point blocked
                    continue;

                // reflection color
                Eigen::Vector3f ref_color =
                    object.brdf(objects[i].light_color, in_dir, out_dir, normal) *
                    fmaxf(normal.dot(in_dir), 0.0f);

                Eigen::Vector3f N = E1.cross(E2);
                Eigen::Vector3f ray = point - sample;
                float N_len = N.norm();
                float ray_len = ray.norm();

                color += ref_color *
                    (N / N_len).dot(ray / ray_len) / (ray_len * ray_len) * (N_len / 2);
            }

    // randomly stop bouncing
    if (config.rng() > config.bounce_chance) return color;

    // try hitting objects

    // incident angle
    float ph = config.rng() * 2.0f * (float)EIGEN_PI;
    float h = config.rng();

    Eigen::Vector3f left = normal.cross(up);
    // reversed incident direction
    Eigen::Vector3f in_dir =
        (up * cosf(ph) + left * sinf(ph)) * sqrtf(1.0f - h * h) +
        normal * h;

    // shoot at reversed incident direction
    float dist;
    int new_object_idx, new_face_idx;
    Eigen::Vector3f sky_color;
    scene.intersect(point, in_dir, dist, new_object_idx, new_face_idx, sky_color);

    Eigen::Vector3f in_color = Eigen::Vector3f::Zero();

    if (new_object_idx < 0) // hit nothing
    {
        in_color = sky_color;
    }
    else if (objects[new_object_idx].type == 1) // hit object
    {
        // hit position
        Eigen::Vector3f new_point = point + in_dir * dist;

        // incident color
        in_color = shade(
            scene, new_object_idx, new_face_idx, new_point, -in_dir, config);
    }

    // reflection color
    Eigen::Vector3f ref_color =
        object.brdf(in_color, in_dir, out_dir, normal) *
        fmaxf(normal.dot(in_dir), 0.0f);

    color += ref_color * (2.0f * (float)EIGEN_PI) / config.bounce_chance;

    return color;
}

Eigen::Vector3f shade_pixel(
    const Scene & scene,
    const Eigen::Vector3f & cam_pos,
    const Eigen::Vector3f & cam_dir,
    const Eigen::Vector3f & cam_up,
    float cam_dist,
    float pixel_x,
    float pixel_y,
    float pixel_width,
    float pixel_height,
    const RenderConfig & config)
{
    Eigen::Vector3f cam_right = cam_dir.cross(cam_up);

    Eigen::Vector3f color = Eigen::Vector3f::Zero();

    for (int i = 0; i < config.samples_per_pixel; i++) // shoot N rays
    {
        // randomly sample a position in the pixel
        float sample_x = pixel_x + pixel_width * config.rng();
        float sample_y = pixel_y + pixel_height * config.rng();
        Eigen::Vector3f sample_dir =
            cam_dir * cam_dist +
            cam_right * sample_x +
            cam_up * sample_y;
        sample_dir.normalize();

        // shoot from camera position to sample position
        float dist;
        int object_idx, face_idx;
        Eigen::Vector3f sky_color;
        scene.intersect(cam_pos, sample_dir, dist, object_idx, face_idx, sky_color);

        auto & objects = scene.get_objects();

        Eigen::Vector3f sample_color;

        if (object_idx < 0) // hit nothing
        {
            sample_color = sky_color;
        }
        else if (objects[object_idx].type == 1) // hit object
        {
            Eigen::Vector3f new_point = cam_pos + sample_dir * dist;

            sample_color = shade(
                scene, object_idx, face_idx, new_point, -sample_dir, config);
        }
        else // hit light
        {
            sample_color = objects[object_idx].light_color;
        }

        color += sample_color / config.samples_per_pixel;
    }

    return color;
}
